﻿#region using directives

using System;
using System.Runtime.InteropServices;

#endregion

namespace Client.UIHelper
{
    /// <summary>
    /// The dataprotector class encrypts and decrpyts data.
    /// </summary>
    public class DataProtector
    {
        #region Store enum

        public enum Store
        {
            USE_MACHINE_STORE = 1,
            USE_USER_STORE
        } ;

        #endregion

        private const int CRYPTPROTECT_LOCAL_MACHINE = 0x4;
        private const int CRYPTPROTECT_UI_FORBIDDEN = 0x1;
        private static readonly IntPtr NullPtr = ((IntPtr) ((0)));

        private readonly Store store;

        /// <summary>
        /// Constructor for DataProtector Class
        /// </summary>
        /// <param name="tempStore"></param>
        public DataProtector(Store tempStore)
        {
            store = tempStore;
        }

        [DllImport("Crypt32.dll", SetLastError = true, CharSet = CharSet.Auto)]
        private static extern bool CryptProtectData(ref DATA_BLOB pDataIn, String szDataDescr, ref DATA_BLOB pOptionalEntropy, IntPtr pvReserved, ref CRYPTPROTECT_PROMPTSTRUCT pPromptStruct, int dwFlags, ref DATA_BLOB pDataOut);

        [DllImport("Crypt32.dll", SetLastError = true, CharSet = CharSet.Auto)]
        private static extern bool CryptUnprotectData(ref DATA_BLOB pDataIn, String szDataDescr, ref DATA_BLOB pOptionalEntropy, IntPtr pvReserved, ref CRYPTPROTECT_PROMPTSTRUCT pPromptStruct, int dwFlags, ref DATA_BLOB pDataOut);

        [DllImport("kernel32.dll", CharSet = CharSet.Auto)]
        private static extern unsafe int FormatMessage(int dwFlags, ref IntPtr lpSource, int dwMessageId, int dwLanguageId, ref String lpBuffer, int nSize, IntPtr* Arguments);

        /// <summary>
        /// Encrypts an array of bytes.
        /// </summary>
        /// <param name="plainText">The byte array to be encrypted.</param>
        /// <param name="optionalEntropy">Optionally used.  Not used in this implementation.</param>
        /// <returns>The encrypted byte array</returns>
        public byte[] Encrypt(byte[] plainText, byte[] optionalEntropy)
        {
            bool retVal = false;

            var plainTextBlob = new DATA_BLOB();
            var cipherTextBlob = new DATA_BLOB();
            var entropyBlob = new DATA_BLOB();

            var prompt = new CRYPTPROTECT_PROMPTSTRUCT();
            InitPromptstruct(ref prompt);

            int dwFlags;
            try
            {
                try
                {
                    int bytesSize = plainText.Length;
                    plainTextBlob.pbData = Marshal.AllocHGlobal(bytesSize);
                    if (IntPtr.Zero == plainTextBlob.pbData)
                    {
                        throw new Exception("Unable to allocate plaintext buffer.");
                    }
                    plainTextBlob.cbData = bytesSize;
                    Marshal.Copy(plainText, 0, plainTextBlob.pbData, bytesSize);
                }
                catch (Exception ex)
                {
                    throw new Exception("Exception marshalling data. " + ex.Message);
                }
                if (Store.USE_MACHINE_STORE == store)
                {
//Using the machine store, should be providing entropy.
                    dwFlags = CRYPTPROTECT_LOCAL_MACHINE | CRYPTPROTECT_UI_FORBIDDEN;
                    //Check to see if the entropy is null
                    if (null == optionalEntropy)
                    {
//Allocate something
                        optionalEntropy = new byte[0];
                    }
                    try
                    {
                        int bytesSize = optionalEntropy.Length;
                        entropyBlob.pbData = Marshal.AllocHGlobal(optionalEntropy.Length);
                        ;
                        if (IntPtr.Zero == entropyBlob.pbData)
                        {
                            throw new Exception("Unable to allocate entropy data buffer.");
                        }
                        Marshal.Copy(optionalEntropy, 0, entropyBlob.pbData, bytesSize);
                        entropyBlob.cbData = bytesSize;
                    }
                    catch (Exception ex)
                    {
                        throw new Exception("Exception entropy marshalling data. " + ex.Message);
                    }
                }
                else
                {
//Using the user store
                    dwFlags = CRYPTPROTECT_UI_FORBIDDEN;
                }
                retVal = CryptProtectData(ref plainTextBlob, "", ref entropyBlob, IntPtr.Zero, ref prompt, dwFlags, ref cipherTextBlob);
                if (false == retVal)
                {
                    throw new Exception("Encryption failed. " + GetErrorMessage(Marshal.GetLastWin32Error()));
                }
            }
            catch (Exception ex)
            {
                throw new Exception("Exception encrypting. " + ex.Message);
            }
            var cipherText = new byte[cipherTextBlob.cbData];
            Marshal.Copy(cipherTextBlob.pbData, cipherText, 0, cipherTextBlob.cbData);
            return cipherText;
        }

        /// <summary>
        /// Decrypts an array of bytes
        /// </summary>
        /// <param name="cipherText">The byte array to decrypt.</param>
        /// <param name="optionalEntropy">Optional.  Not used in this implementation.</param>
        /// <returns>The unencrypted byte array.</returns>
        public byte[] Decrypt(byte[] cipherText, byte[] optionalEntropy)
        {
            bool retVal = false;
            var plainTextBlob = new DATA_BLOB();
            var cipherBlob = new DATA_BLOB();
            var prompt = new CRYPTPROTECT_PROMPTSTRUCT();
            InitPromptstruct(ref prompt);
            try
            {
                try
                {
                    int cipherTextSize = cipherText.Length;
                    cipherBlob.pbData = Marshal.AllocHGlobal(cipherTextSize);
                    if (IntPtr.Zero == cipherBlob.pbData)
                    {
                        throw new Exception("Unable to allocate cipherText buffer.");
                    }
                    cipherBlob.cbData = cipherTextSize;
                    Marshal.Copy(cipherText, 0, cipherBlob.pbData, cipherBlob.cbData);
                }
                catch (Exception ex)
                {
                    throw new Exception("Exception marshalling data. " + ex.Message);
                }
                var entropyBlob = new DATA_BLOB();
                int dwFlags;
                if (Store.USE_MACHINE_STORE == store)
                {
//Using the machine store, should be providing entropy.
                    dwFlags = CRYPTPROTECT_LOCAL_MACHINE | CRYPTPROTECT_UI_FORBIDDEN;
                    //Check to see if the entropy is null
                    if (null == optionalEntropy)
                    {
//Allocate something
                        optionalEntropy = new byte[0];
                    }
                    try
                    {
                        int bytesSize = optionalEntropy.Length;
                        entropyBlob.pbData = Marshal.AllocHGlobal(bytesSize);
                        if (IntPtr.Zero == entropyBlob.pbData)
                        {
                            throw new Exception("Unable to allocate entropy buffer.");
                        }
                        entropyBlob.cbData = bytesSize;
                        Marshal.Copy(optionalEntropy, 0, entropyBlob.pbData, bytesSize);
                    }
                    catch (Exception ex)
                    {
                        throw new Exception("Exception entropy marshalling data. " + ex.Message);
                    }
                }
                else
                {
//Using the user store
                    dwFlags = CRYPTPROTECT_UI_FORBIDDEN;
                }
                retVal = CryptUnprotectData(ref cipherBlob, null, ref entropyBlob, IntPtr.Zero, ref prompt, dwFlags, ref plainTextBlob);
                if (false == retVal)
                {
                    throw new Exception("Decryption failed. " + GetErrorMessage(Marshal.GetLastWin32Error()));
                }
                //Free the blob and entropy.
                if (IntPtr.Zero != cipherBlob.pbData)
                {
                    Marshal.FreeHGlobal(cipherBlob.pbData);
                }
                if (IntPtr.Zero != entropyBlob.pbData)
                {
                    Marshal.FreeHGlobal(entropyBlob.pbData);
                }
            }
            catch (Exception ex)
            {
                throw new Exception("Exception decrypting. " + ex.Message);
            }
            var plainText = new byte[plainTextBlob.cbData];
            Marshal.Copy(plainTextBlob.pbData, plainText, 0, plainTextBlob.cbData);
            return plainText;
        }

        private void InitPromptstruct(ref CRYPTPROTECT_PROMPTSTRUCT ps)
        {
            ps.cbSize = Marshal.SizeOf(typeof (CRYPTPROTECT_PROMPTSTRUCT));
            ps.dwPromptFlags = 0;
            ps.hwndApp = NullPtr;
            ps.szPrompt = null;
        }

        private static unsafe String GetErrorMessage(int errorCode)
        {
            int FORMAT_MESSAGE_ALLOCATE_BUFFER = 0x00000100;
            int FORMAT_MESSAGE_IGNORE_INSERTS = 0x00000200;
            int FORMAT_MESSAGE_FROM_SYSTEM = 0x00001000;
            int messageSize = 255;
            String lpMsgBuf = "";
            int dwFlags = FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS;
            var ptrlpSource = new IntPtr();
            var prtArguments = new IntPtr();
            int retVal = FormatMessage(dwFlags, ref ptrlpSource, errorCode, 0, ref lpMsgBuf, messageSize, &prtArguments);
            if (0 == retVal)
            {
                throw new Exception("Failed to format message for error code " + errorCode + ". ");
            }
            return lpMsgBuf;
        }

        #region Nested type: CRYPTPROTECT_PROMPTSTRUCT

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
        internal struct CRYPTPROTECT_PROMPTSTRUCT
        {
            public int cbSize;
            public int dwPromptFlags;
            public IntPtr hwndApp;
            public String szPrompt;
        }

        #endregion

        #region Nested type: DATA_BLOB

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
        internal struct DATA_BLOB
        {
            public int cbData;
            public IntPtr pbData;
        }

        #endregion
    }
}